#include "inference.h"

Mat letterbox(const cv::Mat& source)
{
    int col = source.cols;
    int row = source.rows;
    int _max = MAX(col, row);
    Mat result = Mat::zeros(_max, _max, CV_8UC3);
    source.copyTo(result(Rect(0, 0, col, row)));
    return result;
}

std :: vector<yolo_Data> Create_Data (std :: vector<yolo_Data> Data, Mat output_buffer, float score_threshold, float scale){
    for (int cls = 4; cls < 8; cls++) {
        yolo_Data SingleData;
        for (int i = 0; i < output_buffer.rows; i++) {
            float class_score = output_buffer.at<float>(i, cls);
            // 只有该类别是四个置信度中最大的，且大于阈值
            if (class_score > score_threshold) {

                SingleData.class_scores.push_back(class_score);

                float cx = output_buffer.at<float>(i, 0);
                float cy = output_buffer.at<float>(i, 1);
                float w = output_buffer.at<float>(i, 2);
                float h = output_buffer.at<float>(i, 3);

                // Get the box
                int left = int((cx - 0.5 * w) * scale);
                int top = int((cy - 0.5 * h) * scale);
                int width = int(w * scale);
                int height = int(h * scale);

                // Get the keypoints
                std::vector<float> keypoints;
                Mat kpts = output_buffer.row(i).colRange(8, 18);
                for (int i = 0; i < 5; i++) {
                    float x = kpts.at<float>(0, i * 2 + 0) * scale;
                    float y = kpts.at<float>(0, i * 2 + 1) * scale;

                    keypoints.push_back(x);
                    keypoints.push_back(y);

                }

                SingleData.boxes.push_back(Rect(left, top, width, height));
                SingleData.objects_keypoints.push_back(keypoints);

            }

        }
        SingleData.class_ids = cls - 4;
        Data.push_back(SingleData);

    }
    return Data;
}

void drawBoundingBox(Mat& img, const Rect& box, int class_id, float score, const Scalar& color) {
    rectangle(img, box, color, 2, 8);
    std::string label = labels[class_id] + std::to_string(score).substr(0, 4);
    Size textSize = cv::getTextSize(label, FONT_HERSHEY_SIMPLEX, 0.5, 1, 0);
    Rect textBox(box.tl().x, box.tl().y - 15, textSize.width, textSize.height + 5);
    cv::rectangle(img, textBox, color, FILLED);
    putText(img, label, Point(box.tl().x, box.tl().y - 5), FONT_HERSHEY_SIMPLEX, 0.5, Scalar(255, 255, 255));
}

void draw_kpt(Mat& img, const std::vector<float>& keypoints, const std::vector<Scalar>& colors) {
    for (int i = 0; i < 5; i++) {
        int x = std::clamp(int(keypoints[i * 2 + 0]), 0, img.cols);
        int y = std::clamp(int(keypoints[i * 2 + 1]), 0, img.rows);
        circle(img, Point(x, y), 5, colors[i], 5);
    }
}
